package com.ai.bayesi.core;

import java.util.*;

/**
 * 贝叶斯网络类
 */
public class BayesianNetwork {
    private List<Node> nodes;
    private Map<String, Node> nodeMap;

    public BayesianNetwork() {
        this.nodes = new ArrayList<>();
        this.nodeMap = new HashMap<>();
    }

    public void addNode(Node node) {
        nodes.add(node);
        nodeMap.put(node.getName(), node);
    }

    public Node getNode(String name) {
        return nodeMap.get(name);
    }

    /**
     * 计算联合概率 P(X1=x1, X2=x2, ..., Xn=xn)
     */
    public double jointProbability(Map<String, String> event) {
        double prob = 1.0;
        for (Node node : nodes) {
            String nodeValue = event.get(node.getName());
            if (nodeValue == null || !node.getValues().contains(nodeValue)) {
                throw new IllegalArgumentException("节点 " + node.getName() + " 的值 " + nodeValue + " 无效");
            }

            // 获取父节点的值
            List<String> parentValues = new ArrayList<>();
            for (Node parent : node.getParents()) {
                parentValues.add(event.get(parent.getName()));
            }

            // 计算并累积概率
            prob *= node.getProbability(nodeValue, parentValues);
        }
        return prob;
    }

    /**
     * 执行概率推理，计算 P(query | evidence)
     */
    public double infer(Map<String, String> query, Map<String, String> evidence) {
        Set<String> overlappingKeys = new HashSet<>(query.keySet());
        overlappingKeys.retainAll(evidence.keySet());
        if (!overlappingKeys.isEmpty()) {
            throw new IllegalArgumentException("查询和证据不能包含相同变量: " + overlappingKeys);
        }
        // 构建已知证据（包含查询变量和证据变量）
        Map<String, String> knownEvidence = new HashMap<>();
        if (evidence != null) {
            knownEvidence.putAll(evidence);
        }

        // 收集所有变量
        Set<String> allVariables = new HashSet<>();
        for (Node node : nodes) {
            allVariables.add(node.getName());
        }

        // 收集已知变量（查询 + 证据）
        Set<String> knownVariables = new HashSet<>(knownEvidence.keySet());

        // 确定隐藏变量（需要边缘化的变量）
        List<String> hiddenVariables = new ArrayList<>();
        for (Node node : nodes) {
            String name = node.getName();
            if (!knownVariables.contains(name)) {
                hiddenVariables.add(name);
            }
        }

        // 生成所有可能的隐藏变量组合
        List<List<String>> hiddenValueCombinations = generateAllPossibleValues(hiddenVariables);

        // 计算分子 P(query, evidence)
        double numerator = 0.0;
        for (List<String> hiddenValues : hiddenValueCombinations) {
            // 构建包含查询、证据和隐藏变量的完整事件
            Map<String, String> fullEvent = new HashMap<>(knownEvidence);
            for (int i = 0; i < hiddenVariables.size(); i++) {
                fullEvent.put(hiddenVariables.get(i), hiddenValues.get(i));
            }
            numerator += jointProbability(fullEvent);
        }

        // 计算分母 P(evidence)
        double denominator = 0.0;
        for (List<String> hiddenValues : hiddenValueCombinations) {
            // 构建只包含证据和隐藏变量的事件（不包含查询变量）
            Map<String, String> evidenceEvent = new HashMap<>(evidence);
            for (int i = 0; i < hiddenVariables.size(); i++) {
                evidenceEvent.put(hiddenVariables.get(i), hiddenValues.get(i));
            }
            denominator += jointProbability(evidenceEvent);
        }

        // 返回条件概率
        return denominator > 0 ? numerator / denominator : 0.0;
    }

    /**
     * 生成所有可能的节点值组合
     */
    private List<List<String>> generateAllPossibleValues(List<String> variables) {
        List<List<String>> result = new ArrayList<>();
        result.add(Collections.emptyList());

        for (String var : variables) { // 按列表顺序处理变量
            Node node = getNode(var);
            List<List<String>> newResult = new ArrayList<>();
            for (String value : node.getValues()) {
                for (List<String> assignment : result) {
                    List<String> newAssignment = new ArrayList<>(assignment);
                    newAssignment.add(value);
                    newResult.add(newAssignment);
                }
            }
            result = newResult;
        }
        return result;
    }
}